package com.baijiaoxi.common.utils.jdk;

import javax.crypto.Cipher;
import java.io.*;
import java.math.BigInteger;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.RSAPrivateKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.util.Base64;

public class RSAUtil {
    public static void main(String[] args) throws Exception {
        String plainText = "Hee";

        // 生成密钥对
        KeyPair keyPairA = generateKeyPair();
        KeyPair keyPairB = generateKeyPair();
        // 获取公钥和私钥
        PublicKey publicKeyA = keyPairA.getPublic();
        PrivateKey privateKeyA = keyPairA.getPrivate();

        PublicKey publicKeyB = keyPairB.getPublic();
        PrivateKey privateKeyB = keyPairB.getPrivate();
       /* serializePublicKey(publicKey,"E://rsa//publicKey");
        publicKey=deserializePublicKey( "E://rsa//publicKey");

        serializePrivateKey(privateKey,"E://rsa//privateKey");
        privateKey=deserializePrivateKey( "E://rsa//privateKey");*/

        // A send msg to B 很慢
        //rsa用来加签或 加密别的加密的密钥


        //A端 持有A的私钥和B的公钥
        // 使用公钥B加密
        byte[] encryptedBytes = encrypt(plainText, publicKeyB);
        String encryptedText = Base64.getEncoder().encodeToString(encryptedBytes);
        System.out.println("密文: " + encryptedText);
        // 使用私钥A加签名
        byte[] signatureBytes = sign(encryptedText, privateKeyA);
        String signatureText = Base64.getEncoder().encodeToString(signatureBytes);
        System.out.println("Signature: " + signatureText);



        //B端 持有B的私钥和A的公钥
        // 使用公钥A验证签名
        boolean isVerified = verify(encryptedText, signatureBytes, publicKeyA);
        System.out.println("Signature Verified: " + isVerified);

        // 使用私钥B解密
        byte[] decryptedBytes = decrypt(Base64.getDecoder().decode(encryptedText), privateKeyB);
        String decryptedText = new String(decryptedBytes);
        System.out.println("Decrypted Text: " + decryptedText);






    }




    public static KeyPair generateKeyPair() throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(4096); // 密钥长度为2048位
        return keyPairGenerator.generateKeyPair();
    }

    public static byte[] encrypt(String plainText, PublicKey publicKey) throws Exception {
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        return cipher.doFinal(plainText.getBytes());
    }

    public static byte[] decrypt(byte[] encryptedBytes, PrivateKey privateKey) throws Exception {
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        return cipher.doFinal(encryptedBytes);
    }

    public static byte[] sign(String plainText, PrivateKey privateKey) throws Exception {
        Signature signature = Signature.getInstance("SHA256withRSA");
        signature.initSign(privateKey);
        signature.update(plainText.getBytes());
        return signature.sign();
    }

    public static boolean verify(String plainText, byte[] signatureBytes, PublicKey publicKey) throws Exception {
        Signature signature = Signature.getInstance("SHA256withRSA");
        signature.initVerify(publicKey);
        signature.update(plainText.getBytes());
        return signature.verify(signatureBytes);
    }
    public static void serializePublicKey(PublicKey publicKey, String keyFilePath) throws IOException, NoSuchAlgorithmException, InvalidKeySpecException {
        try (ObjectOutputStream objectOutputStream = new ObjectOutputStream(new BufferedOutputStream(new FileOutputStream(keyFilePath)))) {
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            RSAPublicKeySpec rsaPublicKeySpec = keyFactory.getKeySpec(publicKey, RSAPublicKeySpec.class);
            objectOutputStream.writeObject(rsaPublicKeySpec.getModulus());
            objectOutputStream.writeObject(rsaPublicKeySpec.getPublicExponent());
        }
    }

    public static PublicKey deserializePublicKey(String keyFilePath) throws IOException, NoSuchAlgorithmException, InvalidKeySpecException, ClassNotFoundException {
        try (ObjectInputStream objectInputStream = new ObjectInputStream(new BufferedInputStream(new FileInputStream(keyFilePath)))) {
            BigInteger modulus = (BigInteger)objectInputStream.readObject();
            BigInteger exponent = (BigInteger)objectInputStream.readObject();
            RSAPublicKeySpec rsaPublicKeySpec = new RSAPublicKeySpec(modulus, exponent);
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            return keyFactory.generatePublic(rsaPublicKeySpec);
        }
    }

    public static void serializePrivateKey(PrivateKey privateKey, String keyFilePath) throws IOException, NoSuchAlgorithmException, InvalidKeySpecException {
        try (ObjectOutputStream objectOutputStream = new ObjectOutputStream(new BufferedOutputStream(new FileOutputStream(keyFilePath)))) {
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            RSAPrivateKeySpec rsaPublicKeySpec = keyFactory.getKeySpec(privateKey, RSAPrivateKeySpec.class);
            objectOutputStream.writeObject(rsaPublicKeySpec.getModulus());
            objectOutputStream.writeObject(rsaPublicKeySpec.getPrivateExponent());
        }
    }

    public static PrivateKey deserializePrivateKey(String keyFilePath) throws IOException, NoSuchAlgorithmException, InvalidKeySpecException, ClassNotFoundException {
        try (ObjectInputStream objectInputStream = new ObjectInputStream(new BufferedInputStream(new FileInputStream(keyFilePath)))) {
            BigInteger modulus = (BigInteger)objectInputStream.readObject();
            BigInteger exponent = (BigInteger)objectInputStream.readObject();
            RSAPrivateKeySpec rsaPrivateKeySpec = new RSAPrivateKeySpec(modulus, exponent);
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            return keyFactory.generatePrivate(rsaPrivateKeySpec);
        }
    }

}
